from .container import ComponentContainer
from .hook.dispatcher import HookDispatcher
from .hook.hook import *


class Estimator:

    def __init__(self, container: ComponentContainer):
        self.container = container
        self.dispatcher = container.dispatcher if container.dispatcher is not None else HookDispatcher()

    def train(self, epochs=1, context="cuda"):
        self.container.verify("train")
        self.container.model = self.container.model.to(context)

        self.container.total_epoch = epochs
        self.container.total_iter = len(self.container.dataloader)

        self.dispatcher.run(BeforeRunHook, self.container)
        for epoch in range(epochs):
            self.container.epoch = epoch

            self.dispatcher.run(BeforeEpochHook, self.container)
            for iteration, data in enumerate(self.container.dataloader):
                data_in_context = [item.to(context) for item in data]
                self.container.iter = iteration
                self.container.data = data_in_context

                self.dispatcher.run(BeforeIterHook, self.container)

                self.container.optimizer.zero_grad()
                self.container.data_evaluator(self.container)
                self.container.loss_evaluator(self.container)
                self.container.grad_evaluator(self.container)

                self.dispatcher.run(AfterIterHook, self.container)
            self.dispatcher.run(AfterEpochHook, self.container)
        self.dispatcher.run(AfterRunHook, self.container)

    def val(self, epochs=1, context="cuda"):
        pass
